from StrategyEngine.TradingAlgorithm import TradingAlgorithm
from StrategyEngine.SimulationParameters import SimulationParameters
from StrategyEngine.TradingCalender import GetCalender
from StrategyEngine.TradingEnvironment import TradingEnvironment
# from StrategyEngine.API import order
from sklearn.metrics import accuracy_score, auc, confusion_matrix, f1_score, precision_score, recall_score,roc_curve
from Core.DataView import DataView
import Core.Gadget as Gadget
import datetime
import random
import pandas as pd
from sklearn import tree
import numpy as np
import copy
import pydotplus  # 导入dot插件库
import matplotlib.pyplot as plt
database = DataView.BatchView()
realtimeView = DataView.RealTimeView()


# Define algorithm
def Initialize(api, context):
    print("initialization")
    #api.PlaceOrder("000001.SZ", 50000)
    pass


def HandledData(api, context, data, dt):
    # print("Handled Data " + str(dt))
    pass


def OnDaily(api, context, dt):
    # print("On Daily " + str(dt))
    pass


def OnWeekly(api, context, dt):
    print(" On Weekly " + str(dt))
    pass


def OnMonthly(api, context, dt):

    #filter={}
    #datetime1 = datetime(2000,1,1)
    #a["$gte"] = datetime1
    #a["$lte"] = dt
    #filter["StadDateTime"] = a

    #database.find("Factor","PB",filter)
    print("  --On Monthly-- " + str(dt))
    # ---查询组合---
    portfolio = api.Portfolio()
    # print("Value")
    # print(portfolio.Value)

    # ---查询持仓---
    #position = api.Position("000001.SZ")
    #if position != None:
        #print(position["Equity"])

    # ---调仓列表---
    symbols = []

    # ---手工指定股票---
    #symbols.append({"Symbol": "000001.SZ"})
    #symbols.append({"Symbol": "000002.SZ"})

    data_test = pd.read_csv('D:/StrategyData' + '/' + 'strategy_data_' + str(dt)[0:4] + str(dt)[5:7] + '.csv',
                            encoding='GBK').set_index('Unnamed: 0')
    data_test['date'] = str(dt)[0:4] + str(dt)[5:7]

    referenceDate = dt + datetime.timedelta(days=-210)
    recentMonths = Gadget.GenerateEndDayofMonth(referenceDate, dt)[0:6]
    data_train = pd.DataFrame()
    for time in recentMonths:
        train = pd.read_csv('D:/StrategyData' + '/' + 'strategy_data_' + str(time)[0:4] + str(time)[5:7] + '.csv',
                            encoding='GBK').set_index('Unnamed: 0')
        train['date'] = str(time)[0:4] + str(time)[5:7]
        data_train = pd.concat([data_train, train], axis=0)

    industry_list = [u'采掘', u'传媒', u'电气设备', u'电子', u'房地产', u'纺织服装', u'非银金融', \
                     u'钢铁', u'公共事业', u'国防军工', u'化工', u'机械设备', u'计算机', u'家用电器', \
                     u'建筑材料', u'建筑建材', u'建筑装饰', u'交通运输', u'交运设备', u'金融服务', u'农林牧渔', \
                     u'汽车', u'轻工制造', u'商业贸易', u'食品饮料', u'通信', u'信息服务', u'信息设备', \
                     u'休闲服务', u'医药生物', u'银行', u'有色金属', u'综合', u'公用事业']

    data_train = data_train.dropna()
    data_test = data_test.dropna()

    for i in range(len(industry_list)):
        insdustry = industry_list[i]
        data_train['industry'] = data_train['industry'].replace(insdustry, i)
        data_test['industry'] = data_test['industry'].replace(insdustry, i)

    # train_data = train_data.reindex(range(len(train_data)))
    xlist = ['pricesales', 'GrossProfitMarginTTM', 'ReceivablesTurnoverTTM', 'ROETTM']
    for x in xlist:
        data_train = data_train[data_train[x] != np.inf]
        data_train = data_train[data_train[x] != -np.inf]
        data_test = data_test[data_test[x] != np.inf]
        data_test = data_test[data_test[x] != -np.inf]

    backlist = ['pricecashflow', 'earningnetincome', 'operatingprofit', 'freecashflow']
    for back in backlist:
        data_train[back] = 1.0 / data_train[back]
        data_test[back] = 1.0 / data_test[back]

    data_train['lab'] = 0  # 拼起来归一化，为了之后分开
    data_test['lab'] = 1

    data_all = pd.concat([data_train, data_test], axis=0)
    data_all_copy = copy.deepcopy(data_all)
    data_all = data_all.drop(columns=['return', 'MonthlyExcessReturn', 'industry', 'lab', 'date'])
    data_all = (data_all - data_all.mean()) / data_all.std()
    data_all['industry'] = data_all_copy['industry']
    data_all['lab'] = data_all_copy['lab']
    data_all['return'] = data_all_copy['return']
    data_all['MonthlyExcessReturn'] = data_all_copy['MonthlyExcessReturn']
    data_all['date'] = data_all_copy['date']

    data_train0 = data_all[data_all['lab'] == 0]
    data_test0 = data_all[data_all['lab'] == 1]

    data_train_copy = copy.deepcopy(data_train0)
    data_test_copy = copy.deepcopy(data_test0)

    lim1 = np.percentile(np.array(data_train_copy[['MonthlyExcessReturn']]), 20)
    lim2 = np.percentile(np.array(data_train_copy[['MonthlyExcessReturn']]), 80)

    data_train0['return_tab'] = np.nan
    data_train0['return_tab'][data_train0['MonthlyExcessReturn'] >= lim2] = 1
    data_train0['return_tab'][
        (data_train0['MonthlyExcessReturn'] > lim1) & (data_train0['MonthlyExcessReturn'] < lim2)] = 0
    data_train0['return_tab'][data_train0['MonthlyExcessReturn'] <= lim1] = -1

    data_test0['return_tab'] = np.nan
    data_test0['return_tab'][data_test0['MonthlyExcessReturn'] >= lim2] = 1
    data_test0['return_tab'][
        (data_test0['MonthlyExcessReturn'] > lim1) & (data_test0['MonthlyExcessReturn'] < lim2)] = 0
    data_test0['return_tab'][data_test0['MonthlyExcessReturn'] <= lim1] = -1

    data_train0 = data_train0.drop(columns=['return', 'MonthlyExcessReturn', 'lab', 'date'])
    data_test0 = data_test0.drop(columns=['return', 'MonthlyExcessReturn', 'lab', 'date'])

    train_y = data_train0[['return_tab']]
    train_x = data_train0.drop(columns=['return_tab'])

    test_y = data_test0[['return_tab']]
    test_x = data_test0.drop(columns=['return_tab'])

    # model = tree.DecisionTreeClassifier(min_samples_leaf=1,random_state=0)
    # data_ori = data_all.drop(columns = ['return', 'MonthlyExcessReturn', 'lab', 'date'])
    # Y = data_ori['return_tab']
    # X = data_ori.drop(columns = ['return_tab'])
    # print(pd.DataFrame(Y).apply(pd.value_counts))
    deplist = [5, 8, 10, 15]
    max_dep = []
    f1_list = []
    '''
    for dep in deplist:
        # dep = 5
        # print('dep', dep)
        model = tree.DecisionTreeClassifier(max_depth=dep, min_samples_leaf=1, random_state=0)
        model.fit(train_x, train_y)
        predicted = model.predict(test_x)
        predict = pd.DataFrame(data=predicted, index=test_x.index, columns=['predict'])
        # print ('train_score', model.score(train_x, train_y))
        f1_s = f1_score(test_y, predict, average='weighted')
        f1_list.append(f1_s)
    best_dep = deplist[f1_list.index(max(f1_list))]
    '''
    best_dep = 5
    Model = tree.DecisionTreeClassifier(max_depth=best_dep, min_samples_leaf=1, random_state=0)
    Model.fit(train_x, train_y)
    Predicted = Model.predict(test_x)
    Predict = pd.DataFrame(data=Predicted, index=test_x.index, columns=['predict'])
    selstock = Predict[Predict['predict'] == 1]
    # print (selstock)
    stocklist = selstock.index.tolist()

    dot_data = tree.export_graphviz(Model, out_file=None, feature_names=train_x.columns.tolist(), max_depth=best_dep, filled=True,rounded=True, class_names = ['-1', '0', '1'])  # 将决策树规则生成dot对象
    graph = pydotplus.graph_from_dot_data(dot_data)  # 通过pydotplus将决策树规则解析为图形
    graph.write_pdf('D:/model_result/trees5/' + str(dt)[0:4] + str(dt)[5:7] + '.pdf')  # 将决策树规则保存为PDF文件
    #if len(droplist) != 0:
        #a = select.drop(index = droplist.index.tolist())
    #num = len(select)
    #reward = data.ix[select.index.tolist(),:][['return', 'MonthlyExcessReturn']]
    #backtest.append([file[14:-4], reward.mean()[0], reward.mean()[1], num])
    #print ('done')
    #feature_importance = model.feature_importances_  # 获得指标重要性
    # feature_importance = pd.DataFrame(feature_importance, index = train_x.columns)
    #plt.figure()
    #plt.bar(np.array(train_x.columns), feature_importance)
    #plt.bar(np.arange(feature_importance.shape[0]), feature_importance)
    # plt.bar(np.arange(feature_importance.shape[0]), feature_importance)  # 画出条形图
    #plt.title('feature importance')  # 子网格标题
    #plt.xlabel('features')  # x轴标题
    #plt.ylabel('importance')  # y轴标题
    #plt.suptitle('classification result')  # 图形总标题
    #plt.show()  # 展示图形

    # ---随机选择100只股票进行交易---
    #instruments = context["Instruments"]
    #for i in range(100):
        #rad1 = random.randint(0, len(instruments) - 1)
        #instrument = instruments[rad1]
        #symbol = instrument["Symbol"]
        #symbols.append({"Symbol": symbol})

    listedInstruments = Gadget.FindListedInstrument(database, datetime1=dt, datetime2=dt+datetime.timedelta(days=31))
    listedInstrumentsList = []
    for instrument in listedInstruments:
        listedInstrumentsList.append(instrument["Symbol"])


    print (len(stocklist))
    if len(stocklist) != 0:
    #print (select_stock)
        for symbol in stocklist:
            if symbol in listedInstrumentsList:
                symbols.append({"Symbol": symbol})

        symbols_table = pd.DataFrame(symbols)
        symbols_table.to_csv('D:/model_result/tree_position5/' + str(dt)[0:4] + str(dt)[5:7] + '.csv', encoding='GBK')

    elif len(stocklist) == 0:
        symbols = []

    api.Rebalance(symbols)
    pass


def OnMonthlyBegin(api, context, dt):
    print("  ++On Monthly Begin++ " + str(dt))
    pass



def Analyze(api, context, results):
    print("Analyze")
    print(results.head())
    results.to_csv('d:/performances.csv')

    import matplotlib.pyplot as plt
    # Plot the Portfolio Return and Excess Return data.
    #
    ax1 = plt.subplot(211)
    results.UnitNetValue.plot(ax=ax1)
    results.Benchmark.plot(ax=ax1)
    ax1.set_ylabel('Net Unit Value')
    #
    ax2 = plt.subplot(212, sharex=ax1)
    results.CumExcessReturn.plot(ax=ax2)
    ax2.set_ylabel('Cummulative Excess Return')

    # Show the plot.
    # plt.gcf().set_size_inches(18, 8)
    plt.show()
    pass


# ---交易日历---
tradingCalender = GetCalender("SH")

#
datetime1 = datetime.datetime(2010, 7, 1)
datetime2 = datetime.datetime(2018, 10, 1)

# ---回测参数设置---
simulatorParameters = SimulationParameters(datetime1=datetime1,
                                           datetime2=datetime2,
                                           trading_calendar=tradingCalender,
                                           data_frequency="monthly")
# ---交易环境：市场数据，业绩基准---
tradingEnvironment = TradingEnvironment(benchmark_symbol="000300.SH",
                                        database=DataView.BatchView(),
                                        realtimeView=DataView.RealTimeView(),
                                        trading_calendar=tradingCalender)

# ---构建策略---
strategy = TradingAlgorithm(name="StrategyTree",
                            initialize=Initialize,
                            handle_data=HandledData,
                            on_daily=OnDaily,
                            on_weekly=OnWeekly,
                            on_monthly=OnMonthly,
                            on_monthly_begin=OnMonthlyBegin,
                            analyze=Analyze,
                            simulator_parameters=simulatorParameters,
                            trading_environment=tradingEnvironment)

# ---策略中需要用到的参数---
context = {}


#instruments = Gadget.FindListedInstrument(database, datetime1, datetime2)
#context["Instruments"] = instruments

# ---开始回测---
statistics = strategy.Run(context=context)

# ---Print Statistics of Strategy---
# statistics.to_csv('myoutput.csv')


